import torch
import torch.nn as nn

cfgs = {
    'A': [64, 'M', 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'],
    'B': [64, 64, 'M', 128, 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'],
    'D': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M', 512, 512, 512, 'M'],
    'E': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 256, 'M', 512, 512, 512, 512, 'M', 512, 512, 512, 512, 'M']
}


def make_layers(cfg, batch_norm=False):
    layers = []

    input_channel = 3
    for l in cfg:
        if l == 'M':
            layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
            continue

        layers += [nn.Conv2d(input_channel, l, kernel_size=3, padding=1)]

        if batch_norm:
            layers += [nn.BatchNorm2d(l)]

        layers += [nn.ReLU(inplace=True)]
        input_channel = l

    return nn.Sequential(*layers)


class VGGNetModel(nn.Module):

    def __init__(self, features, num_class=1000, dropout=0.5):
        super(VGGNetModel, self).__init__()
        self.features = features

        self.classifier = nn.Sequential(

            nn.Linear(512 * 7 * 7, 4096),
            nn.ReLU(inplace=True),
            # [B, 4096]
            nn.Dropout(p=dropout),

            nn.Linear(4096, 4096),
            nn.ReLU(inplace=True),
            # [B, 4096]
            nn.Dropout(p=dropout),

            nn.Linear(4096, num_class),
            # [B, num_class]
        )

    def forward(self, x):
        # [B, 3, 224, 224]
        x = self.features(x)
        # [B, 256, 6, 6]
        x = torch.flatten(x, 1)
        # [B, 256*6*6]
        # 保持 batch 维，只展平特征部分
        x = self.classifier(x)
        # [B, num_class]
        return x


def vgg11(num_class=1000):
    return VGGNetModel(make_layers(cfgs["A"], batch_norm=True), num_class=num_class)


def vgg13(num_class=1000):
    return VGGNetModel(make_layers(cfgs["B"], batch_norm=True), num_class=num_class)


def vgg16(num_class=1000):
    return VGGNetModel(make_layers(cfgs["D"], batch_norm=True), num_class=num_class)


def vgg19(num_class=1000):
    return VGGNetModel(make_layers(cfgs["E"], batch_norm=True), num_class=num_class)
